use std::ptr::null_mut;

/// 该项目下，最大能够申请的内存空间 4KB * 8
const HEAP_MAX_SIZE: usize = (4096 * 8) as usize;
/// 页大小 4KB
const PAGE_SIZE: usize = 4096 as usize;

/// 对参数($val)进行字节对齐
macro_rules! align {
    ($val: expr, $order: expr) => {{
        let o = (1 << $order) - 1;
        ($val + o) & !o
    }};
}

/// 块的分配状态
#[repr(u32)]
enum PayLoadFlags {
    /// 未分配
    Free = 0,
    /// 已分配
    Allocated = 1,
}

/// 单个块的头部信息，包含：块大小、块分配状态
struct PayLoad(u32);

impl PayLoad {
    /// 获取块的大小
    pub fn get_size(&self) -> u32 {
        self.0 & 0xFFFF_FFF8 // 与0xFFFF_FFF8为了实现清除分配信息以及实现8字节对齐
    }

    /// 设置单个块的大小
    pub fn set_size(&mut self, size: u32) {
        self.0 &= 0x0000_0007;        // 先把原来的块大小清除
        self.0 |= size & 0xFFFF_FFF8; // 然后设置新的块大小的值
    }

    /// 获取块的分配状态
    pub fn get_flags(&self) -> u32 {
        self.0 & PayLoadFlags::Allocated as u32
    }

    /// 设置块的分配状态
    pub fn set_flags(&mut self, flag: PayLoadFlags) {
        self.0 &= 0xFFFF_FFF8;
        self.0 |= flag as u32;
    }

    /// 输出块自身的信息
    pub fn print_self(&self) {
        println!(
            "PayLoad: size is {}, flags is {}",
            self.get_size(),
            if self.get_flags() == PayLoadFlags::Allocated as u32 {
                "Allocated"
            } else {
                "Free"
            }
        );
    }
}

/// 堆区信息
struct Allocator {
    /// 模拟堆区
    heap: Box<[u8; HEAP_MAX_SIZE]>,
    /// 堆的起始地址
    heap_start: usize,
    /// 堆的大小
    heap_size: usize,
}

/// 实现堆的操作
impl Allocator {
    /// 初始化堆区的基本信息，设置堆区第一个Free块的参数
    fn new() -> Self {
        let mut allocator = Allocator {
            heap: Box::new([0 as u8; HEAP_MAX_SIZE]),
            heap_start: 0,
            heap_size: PAGE_SIZE,
        };
        unsafe {
            let payload = (&allocator.heap[4]) as *const u8 as *mut PayLoad;
            allocator.heap_start = &allocator.heap[0] as *const u8 as usize;
            (*payload).set_size((PAGE_SIZE - 16) as u32);
            (*payload.add((PAGE_SIZE - 12) / 4)).set_size(PAGE_SIZE as u32 - 16);
        }
        allocator
    }

    /// 获取上一个块的Header，没有返回null_mut
    fn get_prev_header(&self, payload: *mut PayLoad) -> *mut PayLoad {
        assert!(!payload.is_null());
        unsafe {
            let addr = payload as usize;
            if addr - 4 <= self.heap_start {
                return null_mut();
            }

            let prev_footer = payload.sub(1);
            let size = (*prev_footer).get_size() as usize;
            let prev = payload.sub(size / 4 + 2);
            prev
        }
    }

    /// 获取下一个块的Header，没有返回null_mut
    fn get_next_header(&self, payload: *mut PayLoad) -> *mut PayLoad {
        assert!(!payload.is_null());
        unsafe {
            let addr = payload as usize;
            let size = (*payload).get_size() as usize;
            if addr + size as usize >= self.heap_start + self.heap_size {
                return null_mut();
            }
            // 加上footer占用的字节(4)
            let next = payload.add(size / 4 + 2);
            next
        }
    }

    /// 通过块的Header获取Footer
    fn get_payload_footer(&self, payload: *mut PayLoad) -> *mut PayLoad {
        assert!(!payload.is_null());
        unsafe {
            let size = (*payload).get_size() as usize;
            let footer = payload.add(size / 4 + 1);
            footer
        }
    }

    /// 返回堆区的第一个块的Header
    fn get_start_payload(&self) -> *mut PayLoad {
        (self.heap_start - 4) as *mut PayLoad
    }

    /// 返回堆区的最后一个块的Footer
    fn get_last_payload(&self) -> *mut PayLoad {
        let last_footer = (self.heap_start + self.heap_size - 8) as *mut PayLoad;
        unsafe {
            let last_size = (*last_footer).get_size();
            last_footer.sub(last_size as usize / 4 + 1)
        }
    }

    /// 模拟系统调用sys_call，也可以作为接口调用sys_call。失败返回0，成功返回实际大小
    fn sys_call_pages(&self, pages_num: usize) -> usize {
        let post_size = pages_num * 4096;
        post_size
    }

    /// 向系统申请一块地址空间与此连续的页，并返回最后一块的地址。失败返回null_mut
    fn post_space(&self, header: *mut PayLoad, size: usize) -> *mut PayLoad {
        let post_size = align!(size, 12);
        println!("post_size: need get new page size is {}", post_size);
        if self.heap_size + post_size > HEAP_MAX_SIZE {
            return null_mut();
        }
        let get_size = self.sys_call_pages(post_size / PAGE_SIZE) as u32;
        if get_size != 0 {
            unsafe {
                let old_size = (*header).get_size();
                let old_flag = (*header).get_flags();
                if old_flag == PayLoadFlags::Free as u32 {
                    (*header).set_size(old_size + get_size);
                    (*self.get_payload_footer(header)).set_size(old_size + get_size);
                    (*self.get_payload_footer(header)).set_flags(PayLoadFlags::Free);
                    return header;
                } else {
                    let new_header = self.get_next_header(header);
                    (*new_header).set_size(get_size - 8); // 去除Header以及Footer的字节大小
                    (*new_header).set_flags(PayLoadFlags::Free);
                    (*self.get_payload_footer(new_header)).set_size(get_size - 8);
                    (*self.get_payload_footer(new_header)).set_flags(PayLoadFlags::Free);
                    return new_header;
                }
            }
        }
        null_mut()
    }

    /// 尝试alloc，失败返回null_mut
    fn try_allocate(&self, start: *mut PayLoad, post_size: u32) -> *mut u8 {
        let mut header = start;
        unsafe {
            while (header as usize) < self.heap_start + self.heap_size {
                let payload_size = (*header).get_size();
                let payload_flag = (*header).get_flags();
                if payload_flag == PayLoadFlags::Free as u32 {
                    if payload_size == post_size {
                        // (*header).setSize(post_size); // 因为PayLoad大小已经和请求的大小一致，不需要设置PayLoad的大小
                        (*header).set_flags(PayLoadFlags::Allocated);
                        return header.add(1) as *mut u8; // 返回当前payload的地址
                    } else if payload_size > post_size {
                        (*header).set_size(post_size);
                        (*header).set_flags(PayLoadFlags::Allocated);

                        // 设置footer的信息
                        (*header.add((post_size / 4 + 1) as usize)).set_size(post_size);
                        (*header.add((post_size / 4 + 1) as usize))
                            .set_flags(PayLoadFlags::Allocated);

                        // 处理剩余Header的信息
                        (*header.add((post_size / 4 + 2) as usize))
                            .set_size((payload_size - post_size - 8) as u32);
                        (*header.add((post_size / 4 + 2) as usize)).set_flags(PayLoadFlags::Free);
                        (*self.get_payload_footer(header.add((post_size as usize / 4) + 2)))
                            .set_size((payload_size - post_size - 8) as u32);
                        (*self.get_payload_footer(header.add((post_size as usize / 4) + 2)))
                            .set_flags(PayLoadFlags::Free);

                        println!(
                            "Header offset is 0x{:02X}, address offset is 0x{:02X}",
                            header as usize - self.heap_start,
                            header.add(1) as usize - self.heap_start
                        );

                        return header.add(1) as *mut u8;
                    }
                }
                // 如果没找到，跳到下一个PayLoad
                // PayLoad的指针加1，实际添加字节4
                header = header.add((payload_size / 4) as usize + 2);
            }
        }
        null_mut()
    }

    /// alloc一块内存，失败会申请页，再失败返回null_mut
    pub fn alloc(&self, alloc_size: usize) -> *mut u8 {
        assert!(alloc_size < self.heap_size - 8); // 可申请内存大小为总heap大小减去header以及为了字节对齐隔开的4字节(4 + 4 = 8)
        let post_size = align!(alloc_size as u32, 3); // 使其进行8字节对齐
        println!("Alloc size is {:?}", post_size);
        let mut header = self.get_start_payload();
        let mut result = self.try_allocate(header, post_size);

        if result.is_null() {
            // 此时header已经到达了目前堆的末尾，需要连接应该寻找上一个块起始地址
            header = self.post_space(self.get_last_payload(), post_size as usize);
            if !header.is_null() {
                print!(
                    "PayLoad offset is 0x{:02X}: ",
                    header as usize - self.heap_start
                );
                unsafe { (*header).print_self() }
                result = self.try_allocate(header, post_size);
            }
        }

        // 如果申请失败，返回空
        if result.is_null() {
            println!("Alloc failed, return null_mut()");
        }
        result
    }

    /// 释放内存，并合并与之相连的Free块
    pub fn free(&self, target: *mut u8) {
        let addr = align!(target as usize, 3);
        // println!("Free offset address(no align)  is 1x{:02X}", target as usize - self.heap_start);
        println!(
            "Free offset address(has align) is 0x{:02X}",
            addr as usize - self.heap_start
        );
        // 保证释放的地址是落在Allocator管理的范围
        assert!(self.heap_start <= addr && addr <= self.heap_start + self.heap_size);
        unsafe {
            // TODO: 检查二次释放 ????
            let header = (target as *mut PayLoad).sub(1); // 通过block的起始地址，获取header
            assert!((*header).get_flags() == PayLoadFlags::Allocated as u32);

            // 进行释放前的数据准备
            // 释放的可能性有以下4种
            // *A[A->F]A*
            // AF[A->F]A*
            // *A[A->F]FA
            // AF[A->F]FA

            let target_size = (*header).get_size();

            let prev_header = self.get_prev_header(header);
            let next_header = self.get_next_header(header);

            if !prev_header.is_null() {
                print!("Prev 0x{:02X}: ", prev_header as usize - self.heap_start);
                (*prev_header).print_self();
            }
            print!("Head 0x{:02X}: ", header as usize - self.heap_start);
            (*header).print_self();
            if !next_header.is_null() {
                print!("Next 0x{:02X}: ", next_header as usize - self.heap_start);
                (*next_header).print_self();
            }

            let prev_size = if prev_header.is_null() {
                0
            } else {
                (*prev_header).get_size()
            };
            let next_size = if next_header.is_null() {
                0
            } else {
                (*next_header).get_size()
            };

            let prev_flags = if prev_header.is_null() {
                PayLoadFlags::Allocated as u32
            } else {
                (*prev_header).get_flags()
            };
            let next_flags = if next_header.is_null() {
                PayLoadFlags::Allocated as u32
            } else {
                (*next_header).get_flags()
            };

            if prev_flags == PayLoadFlags::Allocated as u32
                && next_flags == PayLoadFlags::Allocated as u32
            {
                // TODO: *A[A->F]A*
                // 不需要合并，直接修改为Free
                println!("Type is *A[A->F]A*, not to merge");
                (*header).set_flags(PayLoadFlags::Free);
                (*self.get_payload_footer(header)).set_flags(PayLoadFlags::Free);
            } else if prev_flags == PayLoadFlags::Free as u32
                && next_flags == PayLoadFlags::Allocated as u32
            {
                // TODO: AF[A->F]A*
                println!(
                    "Type is AF[A->F]A*, merge size is {}",
                    prev_size + target_size + 8
                );
                (*header).set_flags(PayLoadFlags::Free);
                (*prev_header).set_size(prev_size + target_size + 8);
                (*self.get_payload_footer(header)).set_size(prev_size + target_size + 8);
                (*self.get_payload_footer(header)).set_flags(PayLoadFlags::Free);
            } else if prev_flags == PayLoadFlags::Allocated as u32
                && next_flags == PayLoadFlags::Free as u32
            {
                // TODO: *A[A->F]FA
                println!(
                    "Type is *A[A->F]FA, merge size is {}",
                    target_size + 8 + next_size
                );
                (*header).set_flags(PayLoadFlags::Free);
                (*header).set_size(target_size + 8 + next_size);
                (*self.get_payload_footer(next_header)).set_size(target_size + 8 + next_size);
            } else {
                // TODO: AF[A->F]FA
                println!(
                    "Type is AF[A->F]FA, merge size is {}",
                    prev_size + 8 + target_size + 8 + next_size
                );
                (*header).set_flags(PayLoadFlags::Free);
                (*prev_header).set_size(prev_size + 8 + target_size + 8 + next_size);
                (*self.get_payload_footer(next_header))
                    .set_size(prev_size + 8 + target_size + 8 + next_size);
            }
        }
    }

    /// 以Hexdump的方式打印相对地址空间[start, end)的值
    pub fn print_self(&self, start: usize, end: usize) {
        let col = 16;
        print!("     ");
        for i in 0..col {
            print!("{:2X} {}", i, if i % 4 == 3 { " " } else { "" });
        }
        println!("");
        let mut j = 0;
        for i in start..end {
            print!(
                "{}{:02X} {}",
                if j % col == 0 {
                    format!("{:04X} ", i / col)
                } else {
                    String::new()
                },
                self.heap[i],
                if j % col == col - 1 {
                    "\n"
                } else {
                    if j % 4 == 3 {
                        " "
                    } else {
                        ""
                    }
                }
            );
            j = j + 1;
        }
    }
}


/// main 函数
fn main() {
    println!("Allocator !!!");

    let allocator = Allocator::new();
    allocator.print_self(0, 128);
    println!("................................");
    allocator.print_self(PAGE_SIZE - 128, PAGE_SIZE);

    println!("--------------------------------");
    let _addr1 = allocator.alloc(6);
    println!("--------------------------------");
    let _addr2 = allocator.alloc(4088 - 16 - 8 - 32);
    println!("--------------------------------");
    let _addr3 = allocator.alloc(32);

    println!("--------------------------------");
    allocator.free(_addr1);
    println!("--------------------------------");
    allocator.free(_addr3);
    println!("--------------------------------");
    allocator.free(_addr2);
    println!("--------------------------------");

    allocator.print_self(0, 128);
    println!("................................");
    allocator.print_self(PAGE_SIZE - 128, PAGE_SIZE + 128);

    let num_usize = 14;
    println!("{}", align!(num_usize as u32, 3));
}
